import sys
import os
import torch
import numpy as np
import random

from utils import write_csv

def custom_test_template(op_name, TEST_CASES, test_fn, ): 

    print("=" * 50, "\n")
    all_cmp_results = []

    for i in range(len(TEST_CASES)): 
        test_case = TEST_CASES[i]
        print(f"[{op_name} Test] case {i} start. \n")
        cmp_results = test_fn(
            test_case
        )
        all_cmp_results.append(cmp_results)
        print(f"[{op_name} Test] case {i} over. \n")
        print("=" * 50, "\n")
    
    for i in range(len(all_cmp_results)): 
        cmp_results = all_cmp_results[i]
        for j in range(len(cmp_results)): 
            cmp_result = cmp_results[j]
            print(f"[case {i} output {j}]: " + cmp_result["print_str"])


def prof_template(op_name, TEST_CASES, prof_fn, csv_header_list): 
    
    print("=" * 50, "\n")
    data_lists = []
    file_path = f"./{op_name}_prof.csv"

    for i in range(len(TEST_CASES)): 

        print(f"[{op_name} Prof] case {i}")
        test_case = TEST_CASES[i]
        times_list = prof_fn(test_case)
        data_list = test_case + times_list
        data_lists.append(data_list)

    print("=" * 50, "\n")
    write_csv(file_path, csv_header_list, data_lists)
    

